import io
import os
import gradio as gr
import base64
from PIL import Image
from llm import glm4v_9B

history = []


def gen_text(image, question):
    try:
        if not question:
            return "请输出对应任务的prompt"
        res = glm4v_9B(image, question=question)
    except Exception as e:
        print("Error", e)
        res = "解析错误"

    image_base64 = image_to_base64(image)
    image_html = f'''
        <div class="image-container">
            <img src="data:image/png;base64,{image_base64}" class="thumbnail"/>
        </div>
        '''
    history.append(["用户上传的图片", image_html])
    history.append(("模型回答", res))
    return history

def image_to_base64(image):
    """将 PIL 图像保存到临时文件并返回文件路径"""
    filename = os.path.basename(image)  # 提取文件名
    pil_image = Image.open(image)
    temp_file = f"./input_image/{filename}"  # 临时文件路径
    buffered = io.BytesIO()
    pil_image.save(temp_file)
    pil_image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def correct_image():
    """纠正票据图片为水平垂直"""
    pass

def clear_history():
    global history
    history = []
    return history

with gr.Blocks(css="styles.css") as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown('## 产品特点')
            gr.Markdown("""
                1.**上传单张图片（发票、合同）**  
                2.**发票抽取-prompt**：这是一张发票，先利用orc识别出图片内容，再给出如下字段：发票号码，开票日期，购买方名称，购买方纳税人识别号，购买方地址、电话，购买方开户行及账号，货物或应税劳务、服务名称，单价，金额，税率，税额，价税合计小写, 并用json格式输出。  
                3.**合同OCR-prompt**：你是精准的OCR工具。严格按照图片部署输出内容，要求表格部分以markdown格式输出。  
            """)
            clear_button = gr.Button("清空历史", elem_id="clear_button")
        with gr.Column(scale=3):
            chatbot_output = gr.Chatbot(label="glm4v-9B提供服务", elem_id="chatbot_output",
                                        elem_classes="output-box")
            with gr.Row():
                image_input = gr.Image(type="filepath",
                                       label="上传图片",
                                       image_mode="RGB",
                                       elem_id="image_input",
                                       elem_classes="input-box")
                text_input = gr.Textbox(label="输入prompt", lines=2, elem_id="text_input", elem_classes="input-box")
            generate_button = gr.Button("生成", elem_id="generate_button", elem_classes="lar-button")
    generate_button.click(gen_text, inputs=(image_input, text_input), outputs=chatbot_output)
    clear_button.click(clear_history, inputs=None, outputs=chatbot_output)


if __name__ == "__main__":
    demo.launch(server_port=9090, server_name="0.0.0.0")
